!-----------------------------------------------------------------------
! This is a 1-D Monte Carlo solver that transportts neutrons across a
! slab. Vacuum boundry conditions on each side of the slab. This problem
! tallies the absorption, transmission and reflection probabilities.
! Aaron Bevill
! Mateusz Monterial
! 10/20/2013
!-----------------------------------------------------------------------


program MCsolve
IMPLICIT NONE
include 'mpif.h'                ! fortran parallel library
! neutron properties
double precision :: x, E, mu                ! position, energy, directional cosine
double precision :: xn, En, mu_n            ! new pos, energy, directional cosine
double precision :: r, theta, phi           ! distance traveled, scattering, polar angles
! problem/material parameters
double precision :: A, alph, beta, D=1      ! mass number, thickness of slab
integer :: nmat                             ! number of materials
integer :: imodel                           ! model iterator
double precision, dimension(:), allocatable :: micro_t, micro_s ! nuclear interaction data per atom
double precision :: Sig_t, Sig_s            ! nuclear interaction data
double precision, dimension(:), allocatable :: unper_dens ! unperturbed atom densities
double precision, dimension(:,:), allocatable :: per_dens ! perturbed atom densities
! random number properties
integer*8 :: irand=0                    ! random number counter
integer*8 :: max_stride=0, stride=150   ! N random numbers reserved per history
double precision :: rand1, rand2        ! random num storage
integer, parameter :: seed = 86456      ! seed for random num gen
! trace data log
integer, parameter :: segsize=500 ! allocated number of segments in a history
integer, parameter :: nbatch = 100 ! number of histories to run before retracing
double precision, dimension(:,:), allocatable :: log_micro_t, log_micro_s
double precision, dimension(segsize) :: log_sig_t, log_sig_s, log_dtc
integer, dimension(segsize) :: log_loss
integer :: iseg=0, nseg, max_seg=0! segment iterator and max tracker
! other variable
double precision :: ref=0, trn=0, abs=0     ! reflected and transition tallies
integer :: iloss                ! loss type (reflec'n/transmission/absorption)
double precision, dimension(:,:), allocatable :: retrace_tallies ! tallies in the perturbed models
integer*8 :: NPS, i                         ! number of particles simulated
character(32) :: NPS_buff, nmat_buff        ! buffers for reading in nps and nmat
double precision :: wgt                     ! weight of the retraced history
double precision :: sensitivity_delta=0.1   ! the perturbation quantity
logical :: retrace_flag=.true.  ! retrace histories iff true
double precision :: t_begin, t_end        ! timers
! MPI variables
integer, parameter :: root=0    ! the root processor
integer :: errcode, procs, rank   ! error, num of procs, current proc, send/recv request
integer :: dest, source ! send destination, recv source
integer, dimension(7) :: send_req_mpi, recv_req_mpi
integer, dimension(MPI_STATUS_SIZE,7) :: send_stat_mpi, recv_stat_mpi ! send/receive status
integer, dimension(MPI_STATUS_SIZE) :: mpi_stat ! mpi status
double precision :: ref_mpi=0, trn_mpi=0, abs_mpi=0 ! total tallies
integer*8 :: max_stride_mpi=0, max_seg_mpi=0          ! stride, segment size
double precision, dimension(:,:), allocatable :: retrace_tallies_mpi ! tallies in the perturbed models
double precision, dimension(:,:), allocatable :: log_micro_t_mpi, log_micro_s_mpi
double precision, dimension(segsize) :: log_sig_t_mpi, log_sig_s_mpi, log_dtc_mpi
integer,  dimension(segsize) :: log_loss_mpi
integer*8 :: nps_begin, nps_end,  model_begin, model_end ! where to begin or end each model
integer :: batch_total, ibatch=0 ! remeinder for nps and batches
integer iprocs, nseg_mpi, iloss_mpi, irem, inps ! processor iterator, max length of segments

! initilize MPI
call MPI_INIT(errcode)
! get the problem size from input line
call getarg(1, nmat_buff) ! get the number of particles
call getarg(2, NPS_buff)  ! get the number of iterations
! convert buffers to integers
read(nmat_buff, '(i10)') nmat
if (nmat == 0) then
    nmat = 100
end if
read(NPS_buff, '(i10)') NPS
if (NPS == 0) then
    NPS = 1e5
end if
! allocate various arrays
allocate(micro_t(nmat))
allocate(micro_s(nmat))
allocate(unper_dens(nmat))
allocate(per_dens(nmat,nmat))
allocate(log_micro_t(nmat,segsize))
allocate(log_micro_s(nmat,segsize))
allocate(retrace_tallies(nmat,3))
! MPI allocatables
allocate(retrace_tallies_mpi(nmat,3))
allocate(log_micro_t_mpi(nmat,segsize))
allocate(log_micro_s_mpi(nmat,segsize))
retrace_tallies=0
! get total number of processors and rank
call MPI_COMM_SIZE(MPI_COMM_WORLD, procs, errcode)
call MPI_COMM_RANK(MPI_COMM_WORLD, rank, errcode)
! load material properties (no energy dependence)
! add the read in ability so these are not changed
A = 200.
alph = ((A-1.)/(A+1.))**2. ! calculare alpha
beta = 1. - alph

! divide the models among the processors
model_begin = ceiling(real(nmat)/real(procs)*real(rank)) + 1
model_end = ceiling(real(nmat)/real(procs)*(real(rank)+1.))
! fill the atom_densities matrix with unperturbed and perturbed densities
call init_atom_densities(nmat, model_begin,model_end, sensitivity_delta, unper_dens, per_dens)
! calculate how nps should be divided
nps_begin = floor(real(NPS)/real(procs))*rank
if (rank<(procs-1)) then
    nps_end = floor(real(NPS)/real(procs))*((rank)+1.) - 1
else
    ! the remeinder goes to the last processor
    nps_end = NPS-1
end if
! calculate total number of batches
batch_total = floor(real(NPS)/real(procs)/real(nbatch))
! get random number generator to start at their correct positions
call srand(seed)
call advance_rng(nps_begin, irand, stride, max_stride)
call MPI_BARRIER(MPI_COMM_WORLD, errcode)
t_begin = MPI_WTIME()
! main do loop to simulate the particles transport through slab
inps=0
do i=nps_begin,nps_end
    ! increment inps coutner
    inps = inps + 1
    ! initilize neutron properties
    x=0.0           ! position
    E=1.0           ! energy
    mu=1.0          ! incident angle to slab
    call advance_rng(i, irand, stride, max_stride)
    iseg=iseg+1
    segment_loop: do iseg=iseg,segsize
        ! sample random nums for distance and absoprtion
        rand1 = rand()
        rand2 = rand()
        irand = irand + 2
        ! calculate new parameters
        call get_micro_xs(E, nmat, micro_t, micro_s)
        Sig_t = dot_product(micro_t, unper_dens)
        Sig_s = dot_product(micro_s, unper_dens)
        r = -1./Sig_t*dlog(rand1)       ! sample distance traveled
        xn = x + r*mu                   ! calculate new position
        ! log the information for later retracing
        if (retrace_flag) then
            log_micro_t(:, iseg) = micro_t
            log_micro_s(:, iseg) = micro_s
            log_sig_t(iseg) = Sig_t
            log_sig_s(iseg) = Sig_s
            log_dtc(iseg) = r
        end if
        ! check the boundry conditions
        if (xn<0.) then                 ! the particle was reflected
            ref = ref+1.
            log_loss(iseg) = 1
            log_dtc(iseg) = (0. - x) / mu ! distance to escape
            exit segment_loop
        else if (xn>D) then             ! the particle was transmited
            trn = trn+1.
            log_loss(iseg) = 2
            log_dtc(iseg) = (D - x) / mu ! distance to escape
            exit segment_loop
        ! check if absorbed
        else if (rand2 * Sig_t > Sig_s) then
            abs = abs+1.
            log_loss(iseg) = 3
            exit segment_loop
        ! neutrons scatters inside the slab
        else
            rand1 = rand()
            rand2 = rand()
            irand = irand + 2
            ! calculate new energy
            En = E - rand1 * beta * E
            ! calculate new direction cosine
            mu_n = 2 * rand2 - 1
            ! save the new parameters
            E = En
            mu = mu_n
            x = xn
            log_loss(iseg) = 0
        end if
    end do segment_loop
    ! retrace the history in batch mode, make sure each process does the same
    ! number of batches
    if (retrace_flag .AND. mod(inps, nbatch) == 0 .AND. ibatch < batch_total-1 &
        & .OR. i==nps_end .AND. retrace_flag) then ! the last batch
        ibatch = ibatch+1 ! increment the number of batches ran
        nseg = iseg ! nsegments in this batch
        max_seg = max(max_seg,nseg)
        do iprocs=1,procs
            if (iprocs<procs) then
                dest =   mod(rank+1,procs)
                if ((rank-1)<0) then
                    source = procs-1
                else
                    source = rank-1
                end if
                ! i_send logs from previous processor
                call MPI_ISEND(log_micro_t, nmat*nseg , MPI_DOUBLE_PRECISION, &
                                & dest, 1, MPI_COMM_WORLD, send_req_mpi(1), errcode)
                call MPI_ISEND(log_micro_s, nmat*nseg , MPI_DOUBLE_PRECISION, &
                                & dest, 2, MPI_COMM_WORLD, send_req_mpi(2), errcode)
                call MPI_ISEND(log_sig_t, nseg , MPI_DOUBLE_PRECISION, &
                                & dest, 3, MPI_COMM_WORLD, send_req_mpi(3), errcode)
                call MPI_ISEND(log_sig_s, nseg , MPI_DOUBLE_PRECISION, &
                                & dest, 4, MPI_COMM_WORLD, send_req_mpi(4), errcode)
                call MPI_ISEND(log_dtc, nseg , MPI_DOUBLE_PRECISION, &
                                & dest, 5, MPI_COMM_WORLD, send_req_mpi(5), errcode)
                call MPI_ISEND(nseg, 1, MPI_INTEGER, &
                                & dest, 6, MPI_COMM_WORLD, send_req_mpi(6), errcode)
                call MPI_ISEND(log_loss, nseg, MPI_INTEGER, &
                                & dest, 7, MPI_COMM_WORLD, send_req_mpi(7), errcode)
                ! i_receive logs to next processor
                call MPI_IRECV(log_micro_t_mpi, nmat*segsize , MPI_DOUBLE_PRECISION, &
                                & source, 1, MPI_COMM_WORLD, recv_req_mpi(1), errcode)
                call MPI_IRECV(log_micro_s_mpi, nmat*segsize , MPI_DOUBLE_PRECISION, &
                                & source, 2, MPI_COMM_WORLD, recv_req_mpi(2), errcode)
                call MPI_IRECV(log_sig_t_mpi, segsize , MPI_DOUBLE_PRECISION, &
                                & source, 3, MPI_COMM_WORLD, recv_req_mpi(3), errcode)
                call MPI_IRECV(log_sig_s_mpi, segsize , MPI_DOUBLE_PRECISION, &
                                & source, 4, MPI_COMM_WORLD, recv_req_mpi(4), errcode)
                call MPI_IRECV(log_dtc_mpi, segsize , MPI_DOUBLE_PRECISION, &
                                & source, 5, MPI_COMM_WORLD, recv_req_mpi(5), errcode)
                call MPI_IRECV(nseg_mpi, 1, MPI_INTEGER, &
                                & source, 6, MPI_COMM_WORLD, recv_req_mpi(6), errcode)
                call MPI_IRECV(log_loss_mpi, segsize, MPI_INTEGER, &
                                & source, 7, MPI_COMM_WORLD, recv_req_mpi(7), errcode)
            end if
            ! retrace the particles
            call retrace(nmat, segsize, nseg, model_begin, model_end, log_loss , &
                & log_micro_t, log_micro_s, log_sig_t, log_sig_s, log_dtc, &
                & per_dens, retrace_tallies)
            ! wait for send/recv to complete before swapping
            if (iprocs<procs) then
                call MPI_WAITALL(7, send_req_mpi, send_stat_mpi, errcode)
                call MPI_WAITALL(7, recv_req_mpi, recv_stat_mpi, errcode)
                log_micro_t = log_micro_t_mpi
                log_micro_s = log_micro_s_mpi
                log_sig_t   = log_sig_t_mpi
                log_sig_s   = log_sig_s_mpi
                log_dtc     = log_dtc_mpi
                nseg =  nseg_mpi
                log_loss = log_loss_mpi
            end if
        end do
        iseg = 0
    end if
end do
! collect maximum segment
CALL MPI_REDUCE(max_seg, max_seg_mpi, 1, MPI_INTEGER, MPI_MAX, root, &
                & MPI_COMM_WORLD, errcode)
CALL MPI_REDUCE(max_stride, max_stride_mpi, 1, MPI_INTEGER, MPI_MAX, root, &
                & MPI_COMM_WORLD, errcode)
! collect all of the tallies
CALL MPI_REDUCE(ref, ref_mpi, 1, MPI_DOUBLE_PRECISION, MPI_SUM, root, &
                & MPI_COMM_WORLD, errcode)
CALL MPI_REDUCE(trn, trn_mpi, 1, MPI_DOUBLE_PRECISION, MPI_SUM, root, &
                & MPI_COMM_WORLD, errcode)
CALL MPI_REDUCE(abs, abs_mpi, 1, MPI_DOUBLE_PRECISION, MPI_SUM, root, &
                & MPI_COMM_WORLD, errcode)
! reduce the retrace tallies
if (retrace_flag) then
CALL MPI_REDUCE(retrace_tallies(:,1), retrace_tallies_mpi(:,1), nmat, &
                & MPI_DOUBLE_PRECISION, MPI_SUM, root, MPI_COMM_WORLD, errcode)
CALL MPI_REDUCE(retrace_tallies(:,2), retrace_tallies_mpi(:,2), nmat, &
                & MPI_DOUBLE_PRECISION, MPI_SUM, root, MPI_COMM_WORLD, errcode)
CALL MPI_REDUCE(retrace_tallies(:,3), retrace_tallies_mpi(:,3), nmat, &
                & MPI_DOUBLE_PRECISION, MPI_SUM, root, MPI_COMM_WORLD, errcode)
end if

! print the result of the tallies
if (rank==root) then
t_end = MPI_WTIME()
write(*,*) 'P(ref)', ref_mpi/real(nps)
write(*,*) 'P(trn)', trn_mpi/real(nps)
write(*,*) 'P(abs)', abs_mpi/real(nps)
write(*,*) 'max segments', max_seg_mpi
write(*,*) 'max stride', max_stride_mpi
write(*,*) 'Time (seconds)', t_end - t_begin
if (retrace_flag) then
    write(*,*) 'Sensitivity checksums'
    write(*,"(ES26.16)") sum(retrace_tallies_mpi(:,1))
    write(*,"(ES26.16)") sum(retrace_tallies_mpi(:,2))
    write(*,"(ES26.16)") sum(retrace_tallies_mpi(:,3))
    if (nmat == 100 .and. NPS == 1e5) then
        call internal_check(nmat, retrace_tallies_mpi, 1e-13)
    end if
    !write(*,*) "Sensitivity estimates"
    !do imodel = 1, nmat
    !    write(*,"(4X,4ES15.7)") &
    !        (retrace_tallies_mpi(imodel, 1) - ref_mpi) / real(NPS) / sensitivity_delta, &
    !        (retrace_tallies_mpi(imodel, 2) - trn_mpi) / real(NPS) / sensitivity_delta, &
    !        (retrace_tallies_mpi(imodel, 3) - abs_mpi) / real(NPS) / sensitivity_delta !, &
           ! (sum(retrace_tallies_mpi(imodel, :)) - ref_mpi - trn_mpi - abs_mpi) &
            !    / real(NPS) / sensitivity_delta
    !end do
end if
end if
call MPI_FINALIZE(errcode) ! finalize the MPI routine
end program MCsolve


subroutine init_atom_densities(nmat, model_begin, model_end, sensitivity_delta, unper_dens, per_dens)
    implicit none
    integer, intent(in) :: nmat
    integer, intent(in) :: model_begin, model_end
    real, intent(in) :: sensitivity_delta
    real, intent(out), dimension(nmat) :: unper_dens
    real, intent(out), dimension(nmat,nmat) :: per_dens
    !
    integer :: imat, imodel, linspace(nmat)

    linspace = (/(imat, imat=1,nmat)/)
    unper_dens = 0.02 * dexp(-0.001 * linspace)
    do imodel=model_begin,model_end
        per_dens(:, imodel) = unper_dens(:)
        per_dens(imodel, imodel) = per_dens(imodel, imodel) * &
            (1. + sensitivity_delta)
    end do
end subroutine

subroutine get_micro_xs(energy, nmat, micro_t, micro_s)
    implicit none
    real, intent(in)  :: energy
    integer, intent(in) :: nmat
    real, dimension(nmat), intent(out) :: micro_t, micro_s
    !
    integer :: linspace(nmat), imat

    linspace = (/(imat, imat=1,nmat)/)
    micro_t = dsin(linspace * dlog(energy)) + 1.1
    micro_s = micro_t * 0.8 * (dsin(4*linspace*dlog(energy)) + 1)
end subroutine get_micro_xs

subroutine advance_rng(ihist, irand, stride, max_stride)
    implicit none
    integer*8, intent(in) :: ihist, stride
    integer*8, intent(inout) :: irand, max_stride
    !
    double precision :: tmp

    if (irand > 0) then
        max_stride = max(max_stride, irand - stride * (ihist - 1))
    end if
    if (irand > stride * ihist) then
        write(*,*) "WARNING: Stride exceeded."
    end if
    do irand = irand, stride * ihist - 1
        tmp = rand()
    end do
end subroutine

subroutine internal_check(nmat, retrace_tallies, eps)
    implicit none
    integer, intent(in) :: nmat
    double precision, dimension(nmat, 3), intent(in) :: retrace_tallies
    double precision, intent(in) :: eps
    !
    double precision :: s
    logical :: wrong_flag=.false.
    character(LEN=*), PARAMETER :: CONSOLE_GREEN=achar(27)//"[1;32m"
    character(LEN=*), PARAMETER :: CONSOLE_RED=achar(27)//"[1;31m"
    character(LEN=*), PARAMETER :: CONSOLE_DEFAULT=achar(27)//"[0m"

    s = sum(retrace_tallies(:,1))
    wrong_flag = wrong_flag .or. dabs(s / 2.2648737035819185E+06 - 1) > eps
    s = sum(retrace_tallies(:,2))
    wrong_flag = wrong_flag .or. dabs(s / 2.5622963911091546E+06 - 1) > eps
    s = sum(retrace_tallies(:,3))
    wrong_flag = wrong_flag .or. dabs(s / 5.1728524632014949E+06 - 1) > eps
    !
    if (wrong_flag) then
        write(*,*) CONSOLE_RED, "INTERNAL CHECK FAIL", CONSOLE_DEFAULT
    else
        write(*,*) CONSOLE_GREEN, "INTERNAL CHECK PASS", CONSOLE_DEFAULT
    end if
end subroutine internal_check

subroutine retrace(nmat, segsize, nseg, model_begin, model_end, log_loss, &
                & log_micro_t, log_micro_s, log_sig_t, log_sig_s, log_dtc, &
                & per_dens, retrace_tallies)
    implicit none
    integer, intent(in) :: nmat, segsize, nseg, model_begin, model_end
    double precision, intent(in), dimension(nmat,segsize) :: log_micro_t, log_micro_s
    integer, intent(in), dimension(segsize) :: log_loss
    double precision, intent(in), dimension(segsize) :: log_sig_t, log_sig_s, log_dtc
    double precision, intent(in), dimension(nmat,nmat) :: per_dens
    double precision, intent(inout), dimension(nmat, 3) :: retrace_tallies
    double precision  :: wgt, Sig_t, Sig_s
    integer :: iseg, imodel, iloss

    wgt = 1
    do imodel=model_begin,model_end
        do iseg = 1, nseg
            Sig_t = dot_product(log_micro_t(:, iseg), per_dens(:, imodel))
            Sig_s = dot_product(log_micro_s(:, iseg), per_dens(:, imodel))
            iloss = log_loss(iseg)
            if (iloss == 0) then
                wgt = wgt * Sig_s / log_sig_s(iseg) &
                    * dexp(log_dtc(iseg) * (log_sig_t(iseg) - Sig_t))
            ! If escaped, correct by prob of escape
            else if (iloss < 3) then
                wgt = wgt * dexp(log_dtc(iseg) * (log_sig_t(iseg) - Sig_t))
                retrace_tallies(imodel, iloss) = &
                    retrace_tallies(imodel, iloss) + wgt
                wgt = 1
            ! If absorbed, use that correction
            else if (iloss == 3) then
                wgt = wgt * (Sig_t-Sig_s) / (log_sig_t(iseg)-log_sig_s(iseg)) &
                    * dexp(log_dtc(iseg) * (log_sig_t(iseg) - Sig_t))
                retrace_tallies(imodel, iloss) = &
                    retrace_tallies(imodel, iloss) + wgt
                wgt = 1
            end if
        end do
    end do
end subroutine retrace
